import cv2
import numpy as np
import gradio as gr
from PIL import Image

def extract_colors_opencv(image):
    # 转换为numpy数组并缩小图像尺寸
    img = np.array(image)
    # 转换为BGR（OpenCV使用BGR）
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    # 转换为二维矩阵（像素点）
    X = np.float32(img.reshape((-1, 3)))  
    # 定义K-means参数
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    K = 5
    attempts = 10
    flags = cv2.KMEANS_PP_CENTERS
    # 应用K-means聚类
    _, labels, centers = cv2.kmeans(X, K, None, criteria, attempts, flags)
    centers = np.uint8(centers)
    # 将颜色从BGR转换为RGB
    centers = centers[:, ::-1]
    # 转换为十六进制颜色
    hex_colors = ['#%02x%02x%02x' % tuple(center) for center in centers]
    # 创建调色板图像
    palette = Image.new('RGB', (300, 60))
    for i, color in enumerate(centers):
        block = Image.new('RGB', (60, 60), tuple(color))
        palette.paste(block, (i*60, 0))
    return palette, '\n'.join(hex_colors)

interface = gr.Interface(
    fn=extract_colors_opencv,
    inputs=gr.Image(type="pil"),
    outputs=["image", "text"],
    title="主色调调色板提取",
    description="上传照片，获取主色调调色板（5个颜色）和每个颜色的十六进制表示形式。"
)

interface.launch()